#ifndef __SSLCONTEXT__
#define __SSLCONTEXT__

#include <iostream>

#include <openssl/ssl.h>
#include <openssl/err.h>

#include <errno.h>
#include <string.h>

#include "GenericException.hpp"
#include "hexdump.hpp"
#include "TCPConnection.hpp"

/**
 * Wraps OpenSSL calls
 */
class SSLWrapper {

	public:

		/* Initialize OpenSSL and return a global context pointer */
		static SSL_CTX* init() {
			// Global system initialization
			SSL_library_init();
			SSL_load_error_strings();

			// Create our context
			const SSL_METHOD* method = SSLv23_method();
			SSL_CTX* globalContext = SSL_CTX_new(method);

			//set sessionId to prevent error messages when connecting twice using certificates to authenticate
			const char* sessionId = "httpsd";
			SSL_CTX_set_session_id_context(globalContext, (unsigned char*) sessionId, strlen(sessionId));

			return globalContext;
		}

		/* Load the certigicate file into a context */
		static void loadCertificate(SSL_CTX* globalContext, const std::string& certfile) {
			// Load our keys and certificates
			if (!SSL_CTX_use_certificate_chain_file(globalContext, certfile.c_str())) {
				GENERICSTRMEXCEPTION("Can't read certificate file " << certfile << ": " << ERR_error_string(ERR_get_error(), NULL));
			}
		}

		/* Load a private keyfile into a context */
		static void loadPrivatekey(SSL_CTX* globalContext, const std::string& keyfile) {
			// Private key file
			if (!SSL_CTX_use_PrivateKey_file(globalContext, keyfile.c_str(), SSL_FILETYPE_PEM)) {
				GENERICSTRMEXCEPTION("Can't read private key file " << keyfile << ": " << ERR_error_string(ERR_get_error(), NULL));
			}
		}

		/** Load a certificate hierarchy (or possibly only root CA certificate) into a context
		 * @param caList string with filename containing concatenated certificates
		 */
		static void loadHierarchy(SSL_CTX* globalContext, const std::string& caList) {
			// Load the CAs we trust
			if (!SSL_CTX_load_verify_locations(globalContext, caList.c_str(), 0)) {
				GENERICSTRMEXCEPTION("Can't read CA list file " << caList << ": " << ERR_error_string(ERR_get_error(), NULL));
			}
		}

		/* Sets the verify mode of the operation, see 'man SSL_CTX_set_verify' for valid "Mode" parameters */
		static void setVerifyMode(SSL_CTX* globalContext, int mode) {
			SSL_CTX_set_verify(globalContext, mode, NULL);
		}

		/* Return a new session context from a global context */
		static SSL* newConnection(SSL_CTX* globalContext) {
			return SSL_new(globalContext);
		}

		/** Connect a connection (file/socket descriptor), accept it and return a communication channel
		 * @param sd file descriptor to the underlying tcp connection
		 * @param client specifies if OpenSSL should act as a client or server
		 */
		static BIO* sessionToChannel(int sd, SSL* sessionContext, bool client) {
			BIO* sbio;
			sbio = BIO_new_socket(sd, BIO_NOCLOSE);
			SSL_set_bio(sessionContext, sbio, sbio);

			if (client) {
				if (SSL_connect(sessionContext) <= 0) {
					GENERICSTRMEXCEPTION("Cannot connect SSL connection: " << ERR_error_string(ERR_get_error(), NULL));
				}
			} else {
				if (SSL_accept(sessionContext) <= 0) {
					GENERICSTRMEXCEPTION("Cannot accept SSL connection: " << ERR_error_string(ERR_get_error(), NULL));
				}
			}

			BIO* io = BIO_new(BIO_f_buffer());

			BIO* sslBio;
			sslBio = BIO_new(BIO_f_ssl());
			BIO_set_ssl(sslBio, sessionContext, BIO_CLOSE);
			BIO_push(io, sslBio);

			return io;
		}

		/* Reads a single line from an encrypted channel */
		static std::string readLineFromChannel(BIO* channel) {
			char buffer[256];
			BIO_gets(channel, buffer, sizeof(buffer) - 1);
			if (ERR_get_error() != SSL_ERROR_NONE) {
				GENERICSTRMEXCEPTION("SSL write error: " << ERR_error_string(ERR_get_error(), NULL));
			}

			return buffer;
		}

		/* Writes a single line to an encrypted channel */
		static void writeLineToChannel(BIO* channel, const std::string& line) {
			if (BIO_puts(channel, line.c_str()) <= 0) {
				GENERICSTRMEXCEPTION("SSL write error: " << ERR_error_string(ERR_get_error(), NULL));
			}
			if (BIO_flush(channel) < 0) {
				GENERICSTRMEXCEPTION("SSL flush error: " << ERR_error_string(ERR_get_error(), NULL));
			}
		}

		/* Closes a Session Context */
		static void closeSession(SSL* sessionContext) {
			SSL_shutdown(sessionContext);
			SSL_free(sessionContext);
		}

		/* Displays the subject of a connected peer session context, if present */
		static std::string sslPeerCertificateToString(SSL* sessionContext) {
			X509* cert = SSL_get_peer_certificate(sessionContext);
			if (!cert) return "No peer certificate present";
			std::string name = cert->name;
			X509_free(cert);
			return name;
		}

		/* Reads data from a channel */
		static unsigned int readDataFromChannel(BIO *channel, char *buffer, unsigned int maxbuffersize) {
			int read;
			read = BIO_read(channel, buffer, maxbuffersize);
			if (read <= 0) GENERICSTRMEXCEPTION("SSL read error: " << ERR_error_string(ERR_get_error(), NULL));
			return read;
		}
		
		/* Writes data from a channel */
		static void writeDataToChannel(BIO *channel, char *buffer, unsigned int length) {
			int write;
			write = BIO_write(channel, buffer, length);
			if (write <= 0) GENERICSTRMEXCEPTION("SSL write error (Req = " << length << ", written = " << write << "): " << ERR_error_string(ERR_get_error(), NULL));
			(void)BIO_flush(channel);
		}

		/* Connect TCP Connection 1 (with associated secure channel Channel1) to TCP Connection 2 (with associated secure channel Channel1)
		 * and protocol everything on the line to "filename" */
		static void connectChannels(TCPConnection &tcp1, BIO *channel1, TCPConnection &tcp2, BIO *channel2, const std::string &filename) {
			FILE *f;

			f = fopen(filename.c_str(), "w");
			if (!f) GENERICSTRMEXCEPTION("Cannot create " << filename);
			{
				std::stringstream Strm;
				Strm << "Transcript of communication between " << tcp1.getRemoteStr() << " and " << tcp2.getRemoteStr() << "\n\n";
				fprintf(f, "%s", Strm.str().c_str());
			}
			
			tcp1.setNonBlocking();
			tcp2.setNonBlocking();
			
			unsigned int nodata = 0;
			while (true) {
				bool haddata = false;
				char buffer[1024];
				if (tcp1.canReadData() || BIO_pending(channel1)) {
					int bytes;
					bytes = readDataFromChannel(channel1, buffer, sizeof(buffer));
					fprintf(f, "A -> B: %d bytes\n", bytes);
					FileHexDump(f, buffer, bytes);
					writeDataToChannel(channel2, buffer, bytes);
					haddata = true;
				}
				if (tcp2.canReadData() || BIO_pending(channel2)) {
					int bytes;
					bytes = readDataFromChannel(channel2, buffer, sizeof(buffer));
					fprintf(f, "B -> A: %d bytes\n", bytes);
					FileHexDump(f, buffer, bytes);
					writeDataToChannel(channel1, buffer, bytes);
					haddata = true;

				}
				if (!haddata) {
					usleep(100 * 1000);
					nodata++;
				} else {
					nodata = 0;
				}

				/* ATTENTION: REALLY MESSY HACK!!
				   this works for many purposes, especially for our network security
				   exercises. Should be replaced by a proper shutdown detection (TODO? anyone?) */
				if (nodata > 20) {
					tcp1.close();
					tcp2.close();
					return;
				}
			}

			fclose(f);
		}
};

#endif
